package com.xiezc.util;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;

/**
 * 128位AES加密;ECB模式;密钥是 用户输入字符串的MD5码.<br>
 * 补码方式是:补码填充的每个字节 的 前4个bit位表示的是 明文块中信息字节的个数, 后4个bit表示明文块中 补码填充的字节个数 <br>
 * 示例: 如果明文块只有14位 则 需要填充两个字节,则这两个字节是 0xe2
 *
 * @author xiezc
 * @date 2016年9月20日 上午11:10:42
 */
public class AESUtil {

    private AES aes;

    public void setAes128(AES aes128) {
        this.aes = aes128;
    }

    /**
     * 获得加密对象
     *
     * @param key  用户的加密密钥
     * @return
     * @author xiezc
     */
    public static AESUtil getAES128(String key) {
        AESUtil aESUtil = new AESUtil();
        try {
            aESUtil.setAes128(AES.getAES128(MessageDigest.getInstance("MD5").digest(key.getBytes())));
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        return aESUtil;
    }

    /**
     * 判断一个字节是否符合补码的情况
     *
     * @param b
     * @return
     * @author xiezc
     */
    private int isFill(byte b) {
        int y = (b >>> 4) & 15;
        int z = b & 15;
        if (z > 0 && y > 0 && (z + y) == 16) {
            return z;
        }
        return 0;
    }

    /**
     * 判断一个明文块中是否包含填充的补码, 同时返回明文块中字节数量
     *
     * @param out
     * @return
     * @author xiezc
     */
    private int getUsedByte(byte[] out) {
        byte x = out[15];
        int y = isFill(x);
        if (y == 0) {
            return 16;
        }
        for (int i = 0; i < y; i++) {
            if (y == isFill(out[15 - i]) && x == out[15 - i]) {
                continue;
            } else {
                return 16;
            }
        }
        return 16 - y;
    }

    /**
     * 解密
     *
     * @param is
     * @param os
     * @author xiezc
     */
    public void decrypt(InputStream is, OutputStream os) {
        try {
            byte[] b = new byte[16];
            byte[] out ;
            int len ;
            while ((len = is.read(b)) != -1) {
                if (len == 16) {
                    out = aes.decrypt(b);
                    os.write(out, 0, getUsedByte(out));
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 不足一个明文块的补码,补码字节的前面4个bit位表示信息字节的个数,后面4位bit表示补码的字节个数
     *
     * @param is
     * @param os
     * @author xiezc
     */
    public void encrypt(InputStream is, OutputStream os) {
        try {
            byte[] b = new byte[16];
            byte[] out ;
            int len ;
            while ((len = is.read(b)) != -1) {
                if (len == 16) {
                    out = aes.encrypt(b);
                    os.write(out);
                } else if (len < 16) {
                    byte x = (byte) ((len << 4) ^ (16 - len));
                    for (int i = 15; i >= len; i--) {
                        b[i] = x;
                    }
                    out = aes.encrypt(b);
                    os.write(out);
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

}

/**
 * AES加密解密算法的实现<br>
 * 参考的文章是:http://www.mamicode.com/info-detail-514466.html
 *
 * @author xiezc
 * @date 2016年8月30日 下午2:45:36
 */
class AES {

    private byte[][] extendedKey;

    private final byte[] rc = new byte[]{0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, (byte) 0x80, 0x1B, 0x36};
    /**
     * S盒
     */
    private final byte[][] sBox = {{99, 124, 119, 123, -14, 107, 111, -59, 48, 1, 103, 43, -2, -41, -85, 118},
            {-54, -126, -55, 125, -6, 89, 71, -16, -83, -44, -94, -81, -100, -92, 114, -64},
            {-73, -3, -109, 38, 54, 63, -9, -52, 52, -91, -27, -15, 113, -40, 49, 21},
            {4, -57, 35, -61, 24, -106, 5, -102, 7, 18, -128, -30, -21, 39, -78, 117},
            {9, -125, 44, 26, 27, 110, 90, -96, 82, 59, -42, -77, 41, -29, 47, -124},
            {83, -47, 0, -19, 32, -4, -79, 91, 106, -53, -66, 57, 74, 76, 88, -49},
            {-48, -17, -86, -5, 67, 77, 51, -123, 69, -7, 2, 127, 80, 60, -97, -88},
            {81, -93, 64, -113, -110, -99, 56, -11, -68, -74, -38, 33, 16, -1, -13, -46},
            {-51, 12, 19, -20, 95, -105, 68, 23, -60, -89, 126, 61, 100, 93, 25, 115},
            {96, -127, 79, -36, 34, 42, -112, -120, 70, -18, -72, 20, -34, 94, 11, -37},
            {-32, 50, 58, 10, 73, 6, 36, 92, -62, -45, -84, 98, -111, -107, -28, 121},
            {-25, -56, 55, 109, -115, -43, 78, -87, 108, 86, -12, -22, 101, 122, -82, 8},
            {-70, 120, 37, 46, 28, -90, -76, -58, -24, -35, 116, 31, 75, -67, -117, -118},
            {112, 62, -75, 102, 72, 3, -10, 14, 97, 53, 87, -71, -122, -63, 29, -98},
            {-31, -8, -104, 17, 105, -39, -114, -108, -101, 30, -121, -23, -50, 85, 40, -33},
            {-116, -95, -119, 13, -65, -26, 66, 104, 65, -103, 45, 15, -80, 84, -69, 22}};

    /**
     * 逆S盒
     */
    private final byte[][] iSbox = {{82, 9, 106, -43, 48, 54, -91, 56, -65, 64, -93, -98, -127, -13, -41, -5},
            {124, -29, 57, -126, -101, 47, -1, -121, 52, -114, 67, 68, -60, -34, -23, -53},
            {84, 123, -108, 50, -90, -62, 35, 61, -18, 76, -107, 11, 66, -6, -61, 78},
            {8, 46, -95, 102, 40, -39, 36, -78, 118, 91, -94, 73, 109, -117, -47, 37},
            {114, -8, -10, 100, -122, 104, -104, 22, -44, -92, 92, -52, 93, 101, -74, -110},
            {108, 112, 72, 80, -3, -19, -71, -38, 94, 21, 70, 87, -89, -115, -99, -124},
            {-112, -40, -85, 0, -116, -68, -45, 10, -9, -28, 88, 5, -72, -77, 69, 6},
            {-48, 44, 30, -113, -54, 63, 15, 2, -63, -81, -67, 3, 1, 19, -118, 107},
            {58, -111, 17, 65, 79, 103, -36, -22, -105, -14, -49, -50, -16, -76, -26, 115},
            {-106, -84, 116, 34, -25, -83, 53, -123, -30, -7, 55, -24, 28, 117, -33, 110},
            {71, -15, 26, 113, 29, 41, -59, -119, 111, -73, 98, 14, -86, 24, -66, 27},
            {-4, 86, 62, 75, -58, -46, 121, 32, -102, -37, -64, -2, 120, -51, 90, -12},
            {31, -35, -88, 51, -120, 7, -57, 49, -79, 18, 16, 89, 39, -128, -20, 95},
            {96, 81, 127, -87, 25, -75, 74, 13, 45, -27, 122, -97, -109, -55, -100, -17},
            {-96, -32, 59, 77, -82, 42, -11, -80, -56, -21, -69, 60, -125, 83, -103, 97},
            {23, 43, 4, 126, -70, 119, -42, 38, -31, 105, 20, 99, 85, 33, 12, 125}};

    /**
     * 存放密钥
     */
    private final byte[] key;

    /**
     * 构造方法, 同时传入密钥
     *
     * @param key
     * @throws NoSuchAlgorithmException
     */
    private AES(byte[] key) {
        if (key.length != 16) {
            throw new RuntimeException("密钥必须为16个字节");
        }
        this.key = key;
        this.extendedKey = getKeyExpansion(key);
    }

    public static AES getAES128(byte[] key) {
        return new AES(key);
    }

    /**
     * 轮密钥加
     *
     * @param in
     * @param nk
     * @return
     * @author xiezc
     */
    private byte[][] addRoundKey(byte[][] in, int nk) {
        byte[][] inBytes = new byte[4][4];
        for (int i = 0; i < 4; i++) {
            inBytes[i] = gModW(in[i], extendedKey[i + nk * 4]);
        }
        return inBytes;
    }

    /**
     * 乘法, AES中用到的乘法都在此方法中实现了,注意:只实现了AES中的使用到的乘法
     *
     * @param x
     * @param y
     * @return
     * @author xiezc
     */
    private byte chen(byte x, int y) {
        byte out = 0;
        RuntimeException rException = null;
        switch (y) {
            case 1:
                out = x;
                break;
            case 2:
                out = chen2(x);
                break;
            case 3:
                out = (byte) (chen2(x) ^ x);
                break;
            case 9:
                out = (byte) (chen2(chen2(chen2(x))) ^ x);
                break;
            case 0x0b:
                out = (byte) (chen2(chen2(chen2(x))) ^ chen2(x) ^ x);
                break;
            case 0x0d:
                out = (byte) (chen2(chen2(chen2(x))) ^ chen2(chen2(x)) ^ x);
                break;
            case 0x0e:
                out = (byte) (chen2(chen2(chen2(x))) ^ chen2(chen2(x)) ^ chen2(x));
                break;
            default:
                rException = new RuntimeException("参数y只能是特定的值(1,2,3,9,0x0b,0x0d,0x0e)");
                break;
        }
        if (rException != null) {
            throw rException;
        }
        return out;
    }

    /**
     * 乘以2的方法
     *
     * @param inByte
     * @return
     * @author xiezc
     */
    private byte chen2(byte inByte) {
        if (inByte < 0) {
            return (byte) ((inByte << 1) ^ 27);
        } else {
            return (byte) (inByte << 1);
        }
    }

    /**
     * 步骤四 : 一个明文块的解密
     *
     * @author xiezc
     */
    public byte[] decrypt(byte[] in) {
        if (in.length != 16) {
            throw new RuntimeException("输入的待加密数据必须是16个字节");
        }

        byte[][] inBytes = getDiffuseByte(in);
        // 轮密钥加
        inBytes = addRoundKey(inBytes, 10);
        for (int k = 9; k >= 0; k--) {
            // 行移位
            inBytes = rShiftRows(inBytes);
            // 字节变换
            for (int i = 0; i < 4; i++) {
                for (int j = 0; j < 4; j++) {
                    inBytes[i][j] = iSBoxConverter(inBytes[i][j]);
                }
            }
            // 轮密钥加
            inBytes = addRoundKey(inBytes, k);
            if (k != 0) {
                // 列混淆
                inBytes = rMixColumns(inBytes);
            }
        }
        byte[] x = getFocusedByte(inBytes);
        return x;
    }

    /**
     * 步骤四 : 一个明文块的加密
     *
     * @author xiezc
     */
    public byte[] encrypt(byte[] in) {
        if (in.length != 16) {
            throw new RuntimeException("输入的待解密数据必须是16个字节");
        }
        byte[][] inBytes = getDiffuseByte(in);
        // 轮密钥加
        inBytes = addRoundKey(inBytes, 0);
        for (int k = 0; k < 10; k++) {
            // 字节变换
            for (int i = 0; i < 4; i++) {
                for (int j = 0; j < 4; j++) {
                    inBytes[i][j] = sBoxConverter(inBytes[i][j]);
                }
            }
            // 行移位
            inBytes = shiftRows(inBytes);
            if (k != 9) {
                // 列混淆
                inBytes = mixColumns(inBytes);
            }
            // 轮密钥加
            inBytes = addRoundKey(inBytes, k + 1);
        }
        byte[] x = getFocusedByte(inBytes);
        return x;
    }

    /**
     * 密钥扩展中的G函数
     *
     * @param b
     * @param nk
     * @return
     * @author xiezc
     */
    private byte[] g(byte[] b, int nk) {
        // 左移一位
        byte[] x = new byte[4];
        x[0] = b[1];
        x[1] = b[2];
        x[2] = b[3];
        x[3] = b[0];
        // s盒变换
        x[0] = sBoxConverter(x[0]);
        x[1] = sBoxConverter(x[1]);
        x[2] = sBoxConverter(x[2]);
        x[3] = sBoxConverter(x[3]);
        // 与RC模加
        x[0] = (byte) (x[0] ^ rc[nk]);
        return x;
    }

    /**
     * 获得从一维数组转换成的二维数组4*4
     *
     * @param in
     * @return
     * @author xiezc
     */
    private byte[][] getDiffuseByte(byte[] in) {
        byte[][] out = new byte[4][4];
        for (int i = 0; i < 16; i++) {
            out[i / 4][i % 4] = in[i];
        }
        return out;
    }

    /**
     * 获得从二维数组变成的一维数组
     *
     * @param in
     * @return
     * @author xiezc
     */
    private byte[] getFocusedByte(byte[][] in) {
        byte[] out = new byte[16];
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                out[4 * i + j] = in[i][j];
            }
        }
        return out;
    }

    /**
     * 获得密钥
     *
     * @return
     * @author xiezc
     */
    public final byte[] getKey() {
        return key;
    }

    /**
     * 密钥扩展,获得11轮的128字节的密钥
     *
     * @return
     * @author xiezc
     */
    private byte[][] getKeyExpansion(byte[] key) {
        byte[][] w = new byte[44][4];
        try {
            byte[][] bss = getDiffuseByte(key);
            w[0] = bss[0];
            w[1] = bss[1];
            w[2] = bss[2];
            w[3] = bss[3];
            for (int i = 4; i < 44; i++) {
                byte[] x = w[i - 1];
                if (i % 4 == 0) {
                    x = g(w[i - 1], i / 4);
                }
                w[i] = gModW(x, w[i - 4]);
            }
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
        return w;
    }

    /**
     * 两个长度为4的数组相加
     *
     * @param gb
     * @param wb
     * @return
     * @author xiezc
     */
    private byte[] gModW(byte[] gb, byte[] wb) {
        byte[] x = new byte[4];
        for (int i = 0; i < 4; i++) {
            x[i] = (byte) (gb[i] ^ wb[i]);
        }
        return x;
    }

    /**
     * 逆S盒变换
     *
     * @param b
     * @return
     * @author xiezc
     */
    private byte iSBoxConverter(byte b) {
        int x = 15 & (b >> 4);
        int y = 15 & b;
        return iSbox[x][y];
    }

    /**
     * 步骤三 : 正向列混合
     *
     * @author xiezc
     */
    private byte[][] mixColumns(byte[][] inBytes) {
        byte[][] outBytes = new byte[4][4];
        for (int i = 0; i < 4; i++) {
            outBytes[0][i] = (byte) (chen(inBytes[0][i], 2) ^ chen(inBytes[1][i], 3) ^ chen(inBytes[2][i], 1)
                    ^ chen(inBytes[3][i], 1));
            outBytes[1][i] = (byte) (chen(inBytes[0][i], 1) ^ chen(inBytes[1][i], 2) ^ chen(inBytes[2][i], 3)
                    ^ chen(inBytes[3][i], 1));
            outBytes[2][i] = (byte) (chen(inBytes[0][i], 1) ^ chen(inBytes[1][i], 1) ^ chen(inBytes[2][i], 2)
                    ^ chen(inBytes[3][i], 3));
            outBytes[3][i] = (byte) (chen(inBytes[0][i], 3) ^ chen(inBytes[1][i], 1) ^ chen(inBytes[2][i], 1)
                    ^ chen(inBytes[3][i], 2));
        }
        return outBytes;
    }

    /**
     * 步骤三 : 逆向列混合
     *
     * @author xiezc
     */
    private byte[][] rMixColumns(byte[][] inBytes) {
        byte[][] outBytes = new byte[4][4];
        for (int i = 0; i < 4; i++) {
            outBytes[0][i] = (byte) (chen(inBytes[0][i], 0x0E) ^ chen(inBytes[1][i], 0x0B) ^ chen(inBytes[2][i], 0x0D)
                    ^ chen(inBytes[3][i], 0x09));
            outBytes[1][i] = (byte) (chen(inBytes[0][i], 0x09) ^ chen(inBytes[1][i], 0x0E) ^ chen(inBytes[2][i], 0x0B)
                    ^ chen(inBytes[3][i], 0x0D));
            outBytes[2][i] = (byte) (chen(inBytes[0][i], 0x0D) ^ chen(inBytes[1][i], 0x09) ^ chen(inBytes[2][i], 0x0E)
                    ^ chen(inBytes[3][i], 0x0B));
            outBytes[3][i] = (byte) (chen(inBytes[0][i], 0x0B) ^ chen(inBytes[1][i], 0x0D) ^ chen(inBytes[2][i], 0x09)
                    ^ chen(inBytes[3][i], 0x0E));
        }
        return outBytes;
    }

    /**
     * 步骤二 : 逆向行移位
     *
     * @param inBytes
     * @return
     * @author xiezc
     */
    private byte[][] rShiftRows(byte[][] inBytes) {
        // 第二行的移位
        byte x = inBytes[1][0];
        inBytes[1][0] = inBytes[1][1];
        inBytes[1][1] = inBytes[1][2];
        inBytes[1][2] = inBytes[1][3];
        inBytes[1][3] = x;
        // 第三行的移位
        x = inBytes[2][0];
        byte z = inBytes[2][1];
        inBytes[2][0] = inBytes[2][2];
        inBytes[2][1] = inBytes[2][3];
        inBytes[2][2] = x;
        inBytes[2][3] = z;
        // 第四行
        x = inBytes[3][0];
        z = inBytes[3][1];
        byte y = inBytes[3][2];
        inBytes[3][0] = inBytes[3][3];
        inBytes[3][1] = x;
        inBytes[3][2] = z;
        inBytes[3][3] = y;
        return inBytes;
    }

    /**
     * S盒变换
     *
     * @param b
     * @return
     * @author xiezc
     */
    private byte sBoxConverter(byte b) {
        int x = 15 & (b >> 4);
        int y = 15 & b;
        return sBox[x][y];
    }

    /**
     * 步骤二 : 正向行移位
     *
     * @author xiezc
     */
    private byte[][] shiftRows(byte[][] inBytes) {
        byte x = 0;
        // 第二行的移位
        x = inBytes[1][3];
        inBytes[1][3] = inBytes[1][2];
        inBytes[1][2] = inBytes[1][1];
        inBytes[1][1] = inBytes[1][0];
        inBytes[1][0] = x;
        // 第三行的移位
        x = inBytes[2][3];
        byte z = inBytes[2][2];
        inBytes[2][3] = inBytes[2][1];
        inBytes[2][2] = inBytes[2][0];
        inBytes[2][1] = x;
        inBytes[2][0] = z;
        // 第四行
        x = inBytes[3][3];
        z = inBytes[3][2];
        byte y = inBytes[3][1];
        inBytes[3][3] = inBytes[3][0];
        inBytes[3][2] = x;
        inBytes[3][1] = z;
        inBytes[3][0] = y;
        return inBytes;
    }
}